package com.hzsparrow.ai.config;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.ApplicationRunner;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Service;

import com.hzsparrow.ai.entity.ModelConfig;
import com.hzsparrow.ai.entity.PlatformSettings;
import com.hzsparrow.ai.service.PlatformSettingsService;
import com.hzsparrow.ai.service.ModelConfigService;

import lombok.extern.slf4j.Slf4j;

import org.springframework.boot.ApplicationArguments;

/**
 * AI模型初始化服务
 * 在应用启动时初始化默认的对话模型、嵌入模型和向量存储
 */
@Slf4j
@Service
public class AIModelInitService implements ApplicationRunner {
    
    @Autowired
    private PlatformSettingsService platformSettingsService;
    
    @Autowired
    private ModelConfigService modelConfigService;
    
    @Autowired
    private ChatModelManager chatModelManager;
    
    @Autowired
    private EmbeddingModelManager embeddingModelManager;
    
    @Autowired
    private VectorStoreManager vectorStoreManager;
    
    // 当前默认模型缓存
    private volatile ChatModel defaultChatModel;
    private volatile EmbeddingModel defaultEmbeddingModel;
    private volatile VectorStore defaultVectorStore;
    
    private volatile ModelConfig currentChatModelConfig;
    private volatile ModelConfig currentEmbeddingModelConfig;
    
    @Override
    public void run(ApplicationArguments args) throws Exception {
        log.info("开始初始化AI模型配置...");
        
        try {
            // 获取平台设置
            PlatformSettings platformSettings = platformSettingsService.getPlatformSettings();
            
            if (platformSettings != null) {
                // 初始化默认对话模型
                if (platformSettings.getDefaultChatModelId() != null) {
                    initializeDefaultChatModel(platformSettings.getDefaultChatModelId());
                }
                
                // 初始化默认嵌入模型
                if (platformSettings.getDefaultEmbeddingModelId() != null) {
                    initializeDefaultEmbeddingModel(platformSettings.getDefaultEmbeddingModelId());
                }
                
                // 初始化默认向量存储
                if (defaultEmbeddingModel != null) {
                    initializeDefaultVectorStore();
                }
            }
            
            log.info("AI模型配置初始化完成");
        } catch (Exception e) {
            log.error("AI模型配置初始化失败", e);
            // 不抛出异常，允许应用继续启动
        }
    }
    
    /**
     * 初始化默认对话模型
     */
    private void initializeDefaultChatModel(Long chatModelId) {
        try {
            ModelConfig chatModelConfig = modelConfigService.getById(chatModelId);
            if (chatModelConfig != null && "chat".equals(chatModelConfig.getModelType())) {
                this.defaultChatModel = chatModelManager.getOrCreateChatModel(chatModelConfig);
                this.currentChatModelConfig = chatModelConfig;
                log.info("初始化默认对话模型成功: {}", chatModelConfig.getAlias());
            }
        } catch (Exception e) {
            log.error("初始化默认对话模型失败: {}", chatModelId, e);
        }
    }
    
    /**
     * 初始化默认嵌入模型
     */
    private void initializeDefaultEmbeddingModel(Long embeddingModelId) {
        try {
            ModelConfig embeddingModelConfig = modelConfigService.getById(embeddingModelId);
            if (embeddingModelConfig != null && "embedding".equals(embeddingModelConfig.getModelType())) {
                this.defaultEmbeddingModel = embeddingModelManager.getOrCreateEmbeddingModel(embeddingModelConfig);
                this.currentEmbeddingModelConfig = embeddingModelConfig;
                log.info("初始化默认嵌入模型成功: {}", embeddingModelConfig.getAlias());
            }
        } catch (Exception e) {
            log.error("初始化默认嵌入模型失败: {}", embeddingModelId, e);
        }
    }
    
    /**
     * 初始化默认向量存储
     */
    private void initializeDefaultVectorStore() {
        try {
            if (defaultEmbeddingModel != null) {
                this.defaultVectorStore = vectorStoreManager.getDefaultVectorStore(defaultEmbeddingModel);
                log.info("初始化默认向量存储成功");
            }
        } catch (Exception e) {
            log.error("初始化默认向量存储失败", e);
        }
    }
    
    /**
     * 更新默认对话模型
     */
    public void updateDefaultChatModel(Long newChatModelId) {
        try {
            ModelConfig newChatModelConfig = modelConfigService.getById(newChatModelId);
            if (newChatModelConfig != null && "chat".equals(newChatModelConfig.getModelType())) {
                // 清除旧模型缓存
                if (currentChatModelConfig != null) {
                    chatModelManager.clearCache(currentChatModelConfig.getId());
                }
                
                // 创建新模型
                this.defaultChatModel = chatModelManager.getOrCreateChatModel(newChatModelConfig);
                this.currentChatModelConfig = newChatModelConfig;
                
                log.info("更新默认对话模型成功: {}", newChatModelConfig.getAlias());
            }
        } catch (Exception e) {
            log.error("更新默认对话模型失败: {}", newChatModelId, e);
        }
    }
    
    /**
     * 更新默认嵌入模型
     */
    public void updateDefaultEmbeddingModel(Long newEmbeddingModelId) {
        try {
            ModelConfig newEmbeddingModelConfig = modelConfigService.getById(newEmbeddingModelId);
            if (newEmbeddingModelConfig != null && "embedding".equals(newEmbeddingModelConfig.getModelType())) {
                // 清除旧模型缓存
                if (currentEmbeddingModelConfig != null) {
                    embeddingModelManager.clearCache(currentEmbeddingModelConfig.getId());
                }
                
                // 创建新模型
                this.defaultEmbeddingModel = embeddingModelManager.getOrCreateEmbeddingModel(newEmbeddingModelConfig);
                this.currentEmbeddingModelConfig = newEmbeddingModelConfig;
                
                // 重新初始化向量存储
                vectorStoreManager.clearAllCache();
                this.defaultVectorStore = vectorStoreManager.getDefaultVectorStore(defaultEmbeddingModel);
                
                log.info("更新默认嵌入模型成功: {}", newEmbeddingModelConfig.getAlias());
            }
        } catch (Exception e) {
            log.error("更新默认嵌入模型失败: {}", newEmbeddingModelId, e);
        }
    }
    
    /**
     * 获取当前默认对话模型
     */
    public ChatModel getDefaultChatModel() {
        return defaultChatModel;
    }
    
    /**
     * 获取当前默认嵌入模型
     */
    public EmbeddingModel getDefaultEmbeddingModel() {
        return defaultEmbeddingModel;
    }
    
    /**
     * 获取当前默认向量存储
     */
    public VectorStore getDefaultVectorStore() {
        return defaultVectorStore;
    }
    
    /**
     * 获取指定知识库的向量存储
     */
    public VectorStore getVectorStore(String indexName, int dimensions) {
        if (defaultEmbeddingModel == null) {
            throw new IllegalStateException("默认嵌入模型未初始化");
        }
        return vectorStoreManager.getOrCreateVectorStore(defaultEmbeddingModel, indexName, dimensions);
    }
    
    /**
     * 处理模型配置更新事件
     * 当模型配置被修改时，清除相关缓存并重新创建模型实例
     * 
     * @param modelConfig 更新后的模型配置
     */
    public void handleModelConfigUpdate(ModelConfig modelConfig) {
        try {
            log.info("处理模型配置更新: {} - {}", modelConfig.getAlias(), modelConfig.getId());
            
            // 清除对应的模型缓存
            if ("chat".equals(modelConfig.getModelType())) {
                chatModelManager.clearCache(modelConfig.getId());
                
                // 如果更新的是当前默认对话模型，则重新初始化
                if (currentChatModelConfig != null && 
                    currentChatModelConfig.getId().equals(modelConfig.getId())) {
                    
                    this.defaultChatModel = chatModelManager.getOrCreateChatModel(modelConfig);
                    this.currentChatModelConfig = modelConfig;
                    log.info("重新加载当前默认对话模型: {}", modelConfig.getAlias());
                }
            } else if ("embedding".equals(modelConfig.getModelType())) {
                embeddingModelManager.clearCache(modelConfig.getId());
                
                // 如果更新的是当前默认嵌入模型，则重新初始化
                if (currentEmbeddingModelConfig != null && 
                    currentEmbeddingModelConfig.getId().equals(modelConfig.getId())) {
                    
                    this.defaultEmbeddingModel = embeddingModelManager.getOrCreateEmbeddingModel(modelConfig);
                    this.currentEmbeddingModelConfig = modelConfig;
                    
                    // 重新初始化向量存储
                    vectorStoreManager.clearAllCache();
                    this.defaultVectorStore = vectorStoreManager.getDefaultVectorStore(defaultEmbeddingModel);
                    
                    log.info("重新加载当前默认嵌入模型: {}", modelConfig.getAlias());
                }
            }
            
            log.info("模型配置更新处理完成: {}", modelConfig.getAlias());
        } catch (Exception e) {
            log.error("处理模型配置更新失败: {} - {}", modelConfig.getAlias(), e.getMessage(), e);
        }
    }
    
    /**
     * 处理模型配置删除事件
     * 当模型配置被删除时，清除相关缓存
     * 
     * @param modelId 被删除的模型ID
     * @param modelType 模型类型
     */
    public void handleModelConfigDelete(Long modelId, String modelType) {
        try {
            log.info("处理模型配置删除: {} - {}", modelId, modelType);
            
            if ("chat".equals(modelType)) {
                chatModelManager.clearCache(modelId);
                
                // 如果删除的是当前默认对话模型，清除默认模型
                if (currentChatModelConfig != null && 
                    currentChatModelConfig.getId().equals(modelId)) {
                    this.defaultChatModel = null;
                    this.currentChatModelConfig = null;
                    log.warn("当前默认对话模型已被删除，需要重新设置默认模型");
                }
            } else if ("embedding".equals(modelType)) {
                embeddingModelManager.clearCache(modelId);
                
                // 如果删除的是当前默认嵌入模型，清除默认模型和向量存储
                if (currentEmbeddingModelConfig != null && 
                    currentEmbeddingModelConfig.getId().equals(modelId)) {
                    this.defaultEmbeddingModel = null;
                    this.currentEmbeddingModelConfig = null;
                    this.defaultVectorStore = null;
                    vectorStoreManager.clearAllCache();
                    log.warn("当前默认嵌入模型已被删除，需要重新设置默认模型");
                }
            }
            
            log.info("模型配置删除处理完成: {}", modelId);
        } catch (Exception e) {
            log.error("处理模型配置删除失败: {} - {}", modelId, e.getMessage(), e);
        }
    }
} 